package com.yww.datastructure.redblacktree;

import static com.yww.datastructure.redblacktree.IntRedBlackTree.Color.BLACK;
import static com.yww.datastructure.redblacktree.IntRedBlackTree.Color.RED;

/**
 * @author yww
 * @description 红黑树：自平衡二叉树的一种实现
 * @since 2024/2/18 10:02
 */
public class IntRedBlackTree {

    private Node root; // 根节点

    // 新增：按照 ALV 树正常新增即可，遇到红红不平衡时进行调整。
    public void put(int key, Object value) {
        Node p = root;
        Node parent = null;
        // 寻找空位
        while (p != null) {
            parent = p;
            if (p.key > key) {
                p = p.left;
            } else if (p.key < key) {
                p = p.right;
            } else { // 相同key，进行更新
                p.value = value;
                return;
            }
        }
        // 找到空位
        Node inserted = new Node(key, value);
        if (parent == null) { // 树空
            root = inserted;
        } else if (parent.key > key) {
            parent.left = inserted;
            inserted.parent = parent; // 维护parent指针
        } else {
            parent.right = inserted;
            inserted.parent = parent; // 维护parent指针
        }

        // ---------------红红不平衡调整---------------
        fixRedRed(inserted);
    }

    // 调整 x 与 parent 的红红不平衡
    void fixRedRed(Node x) {
        // case 1：插入节点为根节点，将根节点变黑
        if (x == root) {
            x.color = BLACK;
            return;
        }
        // case 2：插入节点的父亲若为黑色，树的红黑性质不变，无需调整
        if (isBlack(x.parent)) {
            return;
        }
        // .............. 红红相邻，冲突 ..............
        // case 3：红红冲突，叔叔为红色
        Node parent = x.parent;
        Node uncle = x.uncle();
        Node grandfather = parent.parent;
        if (isRed(uncle)) {
            parent.color = BLACK; // 父变黑
            uncle.color = BLACK; // 叔变黑
            grandfather.color = RED; // 黑节点增多，祖父变红进行协调
            fixRedRed(grandfather); // 祖父调整为红，可能遇到红红冲突，递归处理
            return;
        }
        // case 4：叔叔为黑色，变色无法平衡
        // 1：父亲为左孩子，插入节点也是左孩子，此时即 LL 不平衡
        if (parent.isLeftChild() && x.isLeftChild()) { // LL
            parent.color = BLACK; // 父变黑
            grandfather.color = RED; // 祖父变红
            rightRotate(grandfather); // 旋转祖父，使得达到红黑平衡
        }
        // 2：父亲为左孩子，插入节点是右孩子，此时即 LR 不平衡
        else if (parent.isLeftChild() && !x.isLeftChild()) { // LR
            leftRotate(parent); // 左旋，成为LL情况
            x.color = BLACK; // x成为新父亲，父变黑
            grandfather.color = RED; // 祖父变红
            rightRotate(x); // 旋转祖父，使得达到红黑平衡
        }
        // 3: 父亲为右孩子，插入节点也是右孩子，此时即 RR 不平衡
        else if (!parent.isLeftChild() && !x.isLeftChild()) { // RR
            parent.color = BLACK; // 父变黑
            grandfather.color = RED; // 祖父变红
            leftRotate(parent); // 叔叔少一个黑色，将父旋转上去
        }
        // 4：父亲为右孩子，插入节点是左孩子，此时即 RL 不平衡
        else { // RL
            rightRotate(parent); // 将父亲右旋下去成为RR情况
            x.color = BLACK; // 新父变黑
            grandfather.color = RED; // 祖父变红
            leftRotate(grandfather); // 叔叔少个黑，左旋祖父下去，达到平衡
        }
    }

    // 根据 key 查找节点
    Node find(int key) {
        Node p = root;
        while (p != null) {
            if (p.key > key) {
                p = p.left;
            } else if (key > p.key) {
                p = p.right;
            } else {
                return p;
            }
        }
        return null;
    }

    // 找剩余节点（删除节点没有孩子，剩余节点为null；只有左孩子，剩余节点为左孩子；只有右孩子，剩余节点为右孩子；有左右孩子，剩余节点为后继）
    Node findReplaced(Node deleted) {
        if (deleted.left == null && deleted.right == null) {
            return null;
        }
        if (deleted.left == null) {
            return deleted.right;
        }
        if (deleted.right == null) {
            return deleted.left;
        }
        // 找后继
        Node s = deleted.right;
        while (s.left != null) {
            s = s.left;
        }
        return s;
    }

    // 红黑树删除
    public void remove(int key) {
        Node deleted = find(key);
        if (deleted == null) {
            return;
        }
        doRemove(deleted);
    }

    // 双黑调整，node节点可能是deleted或replaced
    private void fixDoubleBlack(Node x) {
        if (root == x) { // 递归到根节点，就无需调整了
            return;
        }
        Node parent = x.parent;
        Node sibling = x.sibling();
        if (isRed(sibling)) { // case3：兄弟节点为红
            if (x.isLeftChild()) {
                leftRotate(parent); // 调整节点是左孩子，父亲左旋
            } else {
                rightRotate(parent); // 调整节点是右孩子，父亲右旋
            }
            sibling.color = BLACK;
            parent.color = RED;
            fixDoubleBlack(x); // 此时 case3情况变为了case4或case5，再次进行递归调整
            return;
        }
        // case4、case5情况处理
        if (sibling != null) {
            // case4: 兄弟节点是黑色，且两侄子都是黑色
            if (isBlack(sibling.left) && isBlack(sibling.right)) {
                sibling.color = RED; // 兄弟变红
                if (isRed(parent)) { // 父亲是红色
                    parent.color = BLACK; // 将父亲的颜色变为黑色，所以路径黑色达到平衡
                } else { // 父亲是黑色，无法改变父亲颜色，父亲触发双黑操作
                    fixDoubleBlack(parent);
                }
            }
            // case5：兄弟节点黑色，但是有一个侄子是红色
            else {
                // LL
                if (sibling.isLeftChild() && isRed(sibling.left)) { // 兄弟节点是左孩子，且左侄子是红色，LL不平衡
                    rightRotate(parent);
                    sibling.left.color = BLACK; // 兄弟右旋上去了，父亲右下来了，路径少个黑，左侄子变黑
                    sibling.color = parent.color; // 兄弟颜色变为原父亲颜色，不能打破原来的平衡
                    parent.color = BLACK; // 父亲变黑
                }
                // LR
                else if (sibling.isLeftChild() && isRed(sibling.right)) { // 兄弟节点是左孩子，且右侄子是红色，LR不平衡
                    sibling.right.color = parent.color; // 由于左旋兄弟，会改变兄弟右指针，变色需提前做
                    leftRotate(sibling);
                    rightRotate(parent);
                    // sibling.color = BLACK; 无需这一步，兄弟已经是黑色
                    parent.color = BLACK; // 父变黑
                }
                // RR
                else if (!sibling.isLeftChild() && isRed(sibling.right)) {
                    leftRotate(parent); // 父亲左旋下，兄弟左旋上去
                    sibling.right.color = BLACK;  // 父亲左旋下，兄弟左旋上去，路径少个黑，右侄子变黑
                    sibling.color = parent.color; // 兄弟颜色变为原父亲颜色，不能打破原来的平衡
                    parent.color = BLACK; // 父亲变黑

                }
                // RL
                else {
                    sibling.left.color = parent.color;
                    rightRotate(sibling); // 兄弟右旋下，侄子右旋上
                    leftRotate(parent); // 父亲左旋下，新兄弟（侄子左旋上）
                    parent.color = BLACK;
                }
            }
        } else { // 被调整节点没有黑色节点，让其父亲触发双黑
            fixDoubleBlack(parent);
        }
    }

    // 递归删除方法
    public void doRemove(Node deleted) {
        Node replaced = findReplaced(deleted); // 删除只会剩余的节点
        Node parent = deleted.parent; // 删除节点父节点
        if (replaced == null) { // 没有孩子
            // case 1-1：删的是根节点，且没有孩子
            if (deleted == root) {
                root = null; // 把树root置空既可
                return;
            }
            // case 1-2：删的不是根节点，且没有孩子
            else {
                if (isBlack(deleted)) {
                    // 双黑做复杂调整：没有孩子，先调整deleted，再删除
                    fixDoubleBlack(deleted);
                } else {
                    // case 2-1：删除节点是叶子节点，且是红色，不用做任何处理，达到红黑平衡
                }
                if (deleted.isLeftChild()) {
                    parent.left = null; // 如果删除节点是左孩子，让父亲左指针置为空
                } else {
                    parent.right = null; // 如果删除节点是右孩子，让父亲右指针置为空
                }
                deleted.parent = null; // help gc
            }
            // .... 处理 ....
            return;
        }
        if (deleted.left == null || deleted.right == null) {
            // case 1-1：删的是根节点，只有一个孩子
            if (deleted == root) {
                // 李代桃僵，交换剩余孩子节点删除节点的键值
                deleted.value = replaced.value;
                deleted.key = replaced.key;
                // 最后清空删除节点的孩子，即完成删除
                deleted.left = null;
                deleted.right = null;
            }
            // case 1-2：删的不是根节点，只有一个孩子
            else {
                if (deleted.isLeftChild()) {
                    parent.left = replaced; // 如果删除节点是父亲左孩子，父亲左指针指向剩余节点
                } else {
                    parent.right = replaced; // 如果删除节点是父亲右孩子，父亲右指针指向剩余节点
                }
                replaced.parent = parent;
                deleted.left = deleted.right = deleted.parent = null; // help gc
                if (isBlack(deleted) && isBlack(replaced)) {
                    // 双黑做复杂调整：删除节点有孩子，先删除deleted，再调整replace
                    fixDoubleBlack(replaced);
                } else {
                    // case 2-2 删的是黑，剩下的是红，剩下这个红节点变黑（双黑简单情况）
                    replaced.color = BLACK;
                }
            }
            return;
        }
        // 左右孩子都有，处理（李代桃僵）交互删除节点和后继节点的键值
        Object value = deleted.value;
        int key = deleted.key;
        deleted.key = replaced.key;
        deleted.value = replaced.value;
        replaced.key = key;
        replaced.value = value;
        doRemove(replaced); // 通过使用李代桃僵，使情况变为只有一个孩子，或没有孩子
    }


    enum Color {
        RED, BLACK
    }

    private static class Node {
        int key;
        Object value;
        Node left;
        Node right;
        Node parent; // 父节点，便于操作
        Color color = RED; // 节点颜色，默认新节点都是红色

        public Node(int key, Object value) {
            this.key = key;
            this.value = value;
        }

        // 是否是左孩子
        boolean isLeftChild() {
            return parent != null && parent.left == this;
        }

        // 叔叔节点
        Node uncle() {
            if (parent == null || parent.parent == null) { // 节点没有父亲或者是爷爷，就没有叔叔
                return null;
            }
            if (parent.isLeftChild()) { // 父亲是爷爷的左孩子，叔叔即为爷爷的右孩子
                return parent.parent.right;
            } else {
                return parent.parent.left; // 否则叔叔为爷爷的左孩子
            }
        }

        // 兄弟节点
        Node sibling() {
            if (parent == null) {
                return null;
            }
            if (this.isLeftChild()) {
                return parent.right;
            } else {
                return parent.left;
            }
        }
    }

    // 判红、判黑方法
    boolean isRed(Node node) {
        return node != null && node.color == RED;
    }

    boolean isBlack(Node node) {
        return node == null || node.color == BLACK;
    }

    // 右旋，旋转pink节点
    void rightRotate(Node pink) {
        Node parent = pink.parent;
        Node yellow = pink.left;
        Node green = yellow.right;
        // 给 green 换爹
        pink.left = green;
        if (green != null) {
            green.parent = pink;
        }
        // pink 右旋转下去，yellow 右旋转上去
        yellow.right = pink;
        yellow.parent = pink.parent;
        pink.parent = yellow;
        // 旋转后新根的父子关系处理
        if (parent == null) {
            root = yellow;
        } else if (parent.left == pink) {
            parent.left = yellow;
        } else {
            parent.right = yellow;
        }
    }

    // 左旋，旋转pink节点
    void leftRotate(Node pink) {
        Node parent = pink.parent;
        Node yellow = pink.right;
        Node green = yellow.left;
        // 给 green 换爹
        pink.right = green;
        if (green != null) {
            green.parent = pink;
        }
        // pink 左旋转下去，yellow 左旋转上去
        yellow.left = pink;
        yellow.parent = pink.parent;
        pink.parent = yellow;
        // 旋转后新根的父子关系处理
        if (parent == null) {
            root = yellow;
        } else if (parent.right == pink) {
            parent.right = yellow;
        } else {
            parent.left = yellow;
        }
    }
}
